# -*- coding: utf-8 -*-

'''
    知识库工具
    目前运行此脚本会重新生成vectordb,会加载res/knowledge_base目录下的所有文件
'''

import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

import shutil
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from utils.sysutil import get_knowledge_doc_path, get_vector_store_path

# 1、remove old knowledge base if exists

def clear():
    db_path = get_vector_store_path()
    if os.path.exists(db_path):
        shutil.rmtree(db_path)
        print("remove old knowledge base")
    else:
        print("no old knowledge base")

def refresh():
    print("refresh knowledge base")
    
    vectorstore = Chroma(
        collection_name="knowledge_base",
        embedding_function=HuggingFaceEmbeddings(),
        persist_directory=get_vector_store_path()
    )

    # Load all files in the base path using os.walk
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    for root, dirs, files in os.walk(get_knowledge_doc_path()):
        for file in files:
            path = os.path.join(root, file)
            loader = TextLoader(path, encoding="utf-8")
            docs = loader.load()
            splits = text_splitter.split_documents(docs)
            vectorstore.add_documents(splits)
    
    print("Knowledge base refreshed.")

if __name__ == "__main__":
    clear()
    refresh()
